"""
示例 2: 创建简单的平行板电容器

演示如何：
1. 创建新的 PA 阵列
2. 设置电极几何
3. 定义电压
4. 保存 PA 文件
"""

import sys
import os

# 添加项目根目录到路径
project_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
sys.path.insert(0, project_root)

from SIMION.PA import PA


def create_parallel_plate_capacitor():
    """创建平行板电容器"""
    
    print("\n" + "="*60)
    print("创建平行板电容器")
    print("="*60 + "\n")
    
    # 创建新的 PA 阵列
    # 100x100x1 的 2D 平面阵列
    pa = PA(
        nx=100,                    # X 方向 100 个网格点
        ny=100,                    # Y 方向 100 个网格点
        nz=1,                      # Z 方向 1 个网格点（2D）
        symmetry='planar',         # 平面对称
        max_voltage=10000,         # 最大电压 10kV
        field_type='electrostatic' # 静电场
    )
    
    print(f"创建阵列: {pa.nx()} x {pa.ny()} x {pa.nz()}")
    print(f"对称类型: {pa.symmetry()}")
    print(f"最大电压: {pa.max_voltage()} V\n")
    
    # 设置电极几何
    # 左边的板设置为 +5V
    # 右边的板设置为 -5V
    
    print("设置左侧电极板 (X=20, +5V)...")
    left_plate_x = 20
    left_voltage = 5.0
    
    for y in range(0, pa.ny()):
        # point(x, y, z, is_electrode=1表示是电极, potential=电压值)
        pa.point(left_plate_x, y, 0, 1, left_voltage)
    
    print("设置右侧电极板 (X=80, -5V)...")
    right_plate_x = 80
    right_voltage = -5.0
    
    for y in range(0, pa.ny()):
        pa.point(right_plate_x, y, 0, 1, right_voltage)
    
    # 验证设置
    print("\n验证电极设置:")
    is_elec_left, pot_left = pa.point(left_plate_x, 50, 0)
    is_elec_right, pot_right = pa.point(right_plate_x, 50, 0)
    
    print(f"  左板 ({left_plate_x}, 50, 0): {pot_left}V, 电极={is_elec_left}")
    print(f"  右板 ({right_plate_x}, 50, 0): {pot_right}V, 电极={is_elec_right}")
    
    # 保存文件
    output_file = os.path.join(project_root, "SIMION", "examples", "pa", "parallel_plates.pa#")
    print(f"\n保存到: {output_file}")
    pa.save(output_file)
    
    print("\n" + "="*60)
    print("创建完成!")
    print("="*60)
    print(f"\n文件已保存: parallel_plates.pa#")
    print("可以在 SIMION 中打开此文件查看电极几何")
    print("提示: 需要在 SIMION 中 Refine 后才能看到电位分布\n")
    
    return pa


def create_simple_einzel_lens():
    """创建简单的 Einzel 透镜"""
    
    print("\n" + "="*60)
    print("创建 Einzel 透镜")
    print("="*60 + "\n")
    
    # 创建圆柱对称阵列
    pa = PA(
        nx=100,                    # 径向方向
        ny=200,                    # 轴向方向
        nz=1,                      # Z 方向（不使用）
        symmetry='cylindrical',    # 圆柱对称
        max_voltage=1000,          # 最大电压 1kV
        field_type='electrostatic'
    )
    
    print(f"创建阵列: {pa.nx()} x {pa.ny()} x {pa.nz()}")
    print(f"对称类型: {pa.symmetry()}")
    print(f"最大电压: {pa.max_voltage()} V\n")
    
    # 创建三个圆柱电极
    # 外径
    outer_radius = 40
    inner_radius = 10
    
    # 第一个电极 (接地)
    print("设置第一个电极 (Y=30-50, 0V)...")
    for y in range(30, 51):
        for r in range(inner_radius, outer_radius + 1):
            pa.point(r, y, 0, 1, 0.0)
    
    # 第二个电极 (中心电极，施加电压)
    print("设置第二个电极 (Y=90-110, 100V)...")
    for y in range(90, 111):
        for r in range(inner_radius, outer_radius + 1):
            pa.point(r, y, 0, 1, 100.0)
    
    # 第三个电极 (接地)
    print("设置第三个电极 (Y=150-170, 0V)...")
    for y in range(150, 171):
        for r in range(inner_radius, outer_radius + 1):
            pa.point(r, y, 0, 1, 0.0)
    
    # 保存文件
    output_file = os.path.join(project_root, "SIMION", "examples", "pa", "einzel_lens.pa#")
    print(f"\n保存到: {output_file}")
    pa.save(output_file)
    
    print("\n" + "="*60)
    print("创建完成!")
    print("="*60)
    print(f"\n文件已保存: einzel_lens.pa#")
    print("这是一个简单的 Einzel 透镜，可以用于聚焦离子束\n")
    
    return pa


if __name__ == "__main__":
    # 创建平行板电容器
    pa1 = create_parallel_plate_capacitor()
    
    # 创建 Einzel 透镜
    pa2 = create_simple_einzel_lens()
    
    print("\n所有示例创建完成!")
    print("可以在 SIMION 中打开这些 PA 文件:")
    print("  - parallel_plates.pa#")
    print("  - einzel_lens.pa#")

